"""
训练任务数据模型
"""
from datetime import datetime
from typing import Optional
from sqlalchemy import (
    Column, Integer, String, Boolean, DateTime, Text, 
    JSON, ForeignKey, Float, Enum as SQLEnum
)
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
import enum

from app.core.database import Base


class TrainingStatus(str, enum.Enum):
    """训练状态枚举"""
    PENDING = "pending"
    QUEUED = "queued"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class TrainingType(str, enum.Enum):
    """训练类型枚举"""
    SUPERVISED = "supervised"
    UNSUPERVISED = "unsupervised"
    REINFORCEMENT = "reinforcement"
    FEDERATED = "federated"
    TRANSFER = "transfer"
    INCREMENTAL = "incremental"


class TrainingJob(Base):
    """训练任务模型"""
    __tablename__ = "training_jobs"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), nullable=False, index=True)
    description = Column(Text, nullable=True)
    training_type = Column(SQLEnum(TrainingType), nullable=False)
    status = Column(SQLEnum(TrainingStatus), default=TrainingStatus.PENDING, nullable=False)
    
    # 训练配置
    model_architecture = Column(String(50), nullable=False)  # transformer, cnn, rnn等
    hyperparameters = Column(JSON, default=dict, nullable=False)
    training_config = Column(JSON, default=dict, nullable=False)
    resource_requirements = Column(JSON, default=dict, nullable=False)  # GPU, CPU, memory需求
    
    # 数据配置
    dataset_config = Column(JSON, default=dict, nullable=False)
    data_split_config = Column(JSON, default=dict, nullable=False)  # train/val/test分割
    augmentation_config = Column(JSON, default=dict, nullable=True)
    
    # 进度信息
    progress_percentage = Column(Float, default=0.0, nullable=False)
    current_epoch = Column(Integer, default=0, nullable=False)
    total_epochs = Column(Integer, nullable=False)
    
    # 性能指标
    metrics = Column(JSON, default=dict, nullable=False)  # 实时指标
    best_metrics = Column(JSON, default=dict, nullable=False)  # 最佳指标
    
    # 执行信息
    worker_node = Column(String(100), nullable=True)  # 执行节点
    gpu_ids = Column(JSON, default=list, nullable=False)  # 使用的GPU
    log_path = Column(String(255), nullable=True)
    checkpoint_path = Column(String(255), nullable=True)
    
    # 时间信息
    estimated_duration = Column(Integer, nullable=True)  # 预估时长(秒)
    actual_duration = Column(Integer, nullable=True)  # 实际时长(秒)
    started_at = Column(DateTime(timezone=True), nullable=True)
    completed_at = Column(DateTime(timezone=True), nullable=True)
    
    # 错误信息
    error_message = Column(Text, nullable=True)
    error_traceback = Column(Text, nullable=True)
    
    # 关联关系
    project_id = Column(Integer, ForeignKey("projects.id"), nullable=False)
    created_by = Column(Integer, ForeignKey("users.id"), nullable=False)
    parent_job_id = Column(Integer, ForeignKey("training_jobs.id"), nullable=True)  # 支持任务链
    
    # 时间戳
    created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False)
    updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False)
    
    # 关系
    project = relationship("Project", back_populates="training_jobs")
    creator = relationship("User", foreign_keys=[created_by])
    parent_job = relationship("TrainingJob", remote_side=[id])
    child_jobs = relationship("TrainingJob", remote_side=[parent_job_id])
    models = relationship("Model", back_populates="training_job")
    
    def __repr__(self):
        return f"<TrainingJob(id={self.id}, name='{self.name}', status='{self.status}')>"